{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# bitfit : Bias-term Fine-tuning",
   "id": "713e058061119f5f"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# 模型显存占用估算 1.3B模型\n",
    "1. 权重 1.3B * 4B = 5.2G\n",
    "2. 梯度 1.3B * 4B = 5.2G\n",
    "3. 优化器 1.3B * 4B * 2= 10.4G\n",
    "4. total: 20.8G"
   ],
   "id": "2109993c3a52ec45"
  },
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "from datasets import Dataset\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "755038541d1de431",
   "metadata": {},
   "source": "ds = Dataset.load_from_disk(\"data/alpaca_data_zh\")",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "41ac2c304f09f3c8",
   "metadata": {},
   "source": [
    "model_name = \"Langboat/bloom-1b4-zh\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)\n",
    "model.cuda()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "18873d1a865eca39",
   "metadata": {},
   "source": "sum(param.numel() for param in model.parameters())",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def process_function(example):\n",
    "    MAX_LEN = 256\n",
    "    instruction = tokenizer(\"\\n\".join([\"Human:\" + example[\"instruction\"], example[\"input\"]]).strip() + \"\\n\\nAssistant:\")\n",
    "    response = tokenizer(example[\"output\"] + tokenizer.eos_token)\n",
    "    input_ids = instruction[\"input_ids\"] + response[\"input_ids\"]\n",
    "    attention_mask = instruction[\"attention_mask\"] + response[\"attention_mask\"]\n",
    "    labels = [-100] * len(instruction[\"input_ids\"]) + response[\"input_ids\"]\n",
    "    if len(input_ids) > MAX_LEN:\n",
    "        input_ids = input_ids[:MAX_LEN]\n",
    "        attention_mask = attention_mask[:MAX_LEN]\n",
    "        labels = labels[:MAX_LEN]\n",
    "\n",
    "    return {\n",
    "        \"input_ids\": input_ids,\n",
    "        \"attention_mask\": attention_mask,\n",
    "        \"labels\": labels\n",
    "    }\n",
    "\n",
    "remove_columns = ds.column_names\n",
    "print(ds)\n",
    "ds = ds.map(process_function)\n",
    "print(ds)"
   ],
   "id": "944e67ab80b98149",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "冻结非偏置参数",
   "id": "2b4ade00d1461d32"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "for name, param in model.named_parameters():\n",
    "    if \"bias\" not in name:\n",
    "        param.requires_grad = False"
   ],
   "id": "e803a534077c0a54",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "统计待训练参数",
   "id": "12d6aa768fc3502c"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "sum(param.numel() for param in model.parameters() if param.requires_grad)",
   "id": "6878b4beb762ba58",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=\"output\",\n",
    "    per_device_train_batch_size=6,\n",
    "    gradient_accumulation_steps=8,\n",
    "    logging_steps=10,\n",
    "    num_train_epochs=1\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model,\n",
    "    args=args,\n",
    "    train_dataset=ds,\n",
    "    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),\n",
    ")\n",
    "\n",
    "trainer.train()"
   ],
   "id": "9e5ce23455d8f824",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "86774aa17b3ebdb9",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
